package com.timeriver.machine_learning.binaryclassification

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object LogisticRegressionAlgByJDBC {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .appName("读取JDBC数据源进行模型训练")
      .master("local[6]")
      .getOrCreate()

    val df: DataFrame = session.read
      .format("jdbc")
      .option("url", "jdbc:mysql://10.0.24.197:3306/ml_datasets")
      .option("driver", "com.mysql.jdbc.Driver")
      .option("dbtable", "breast_cancer_wisconsin")
      .option("user", "root")
      .option("password", "123456")
      .load()

    df.createOrReplaceTempView("breast_cancer_wisconsin")

    val rawData: DataFrame = session.sql("select clump_thickness,uniformity_of_cell_size,uniformity_of_cell_shape,marginal_adhesion," +
      "single_epithelial_cell_size,bare_nuclei,blan_chromatin,normal_nucleoli,mitoses,class from breast_cancer_wisconsin")

    /** 获取特征列字段数组 */
    val inputCols: Array[String] = "clump_thickness,uniformity_of_cell_size,uniformity_of_cell_shape,marginal_adhesion,single_epithelial_cell_size,bare_nuclei,blan_chromatin,normal_nucleoli,mitoses".split(",")

    /** 过滤缺失值 */
    val value: Dataset[Row] = rawData.filter(!_.anyNull)

    /** 构建特征列向量 */
    val data: DataFrame = new VectorAssembler()
      .setInputCols(inputCols)
      .setOutputCol("features")
      .transform(value)

    val middle_data: DataFrame = data.select("features", "class")

    val Array(trainData, testData) = middle_data.randomSplit(Array(0.7, 0.3), 123L)

    val regression: LogisticRegression = new LogisticRegression()
      .setFeaturesCol("features")
      .setLabelCol("class")
      .setMaxIter(50)
      .setFitIntercept(true)

    val model: LogisticRegressionModel = regression.fit(trainData)

    val pred: DataFrame = model.transform(testData)

    pred.show(5, false)

    session.stop()
  }
}
